from pytorch_lightning.callbacks import ModelCheckpoint
import time
from datam import *
from model import *
import pytorch_lightning as pl
from pytorch_lightning import Trainer
from config import  config,logger



if __name__ == '__main__':

    pl.seed_everything(1234)

    trainer = Trainer(gpus=config.AVAIL_GPUS, max_epochs=config.max_epochs, logger=logger)
    data_mnist = DataM(config.data_dir,config.BATCH_SIZE,config.AVAIL_GPUS)

    model = Model(Backbone())

    #训练模型
    trainer.fit(model,data_mnist)
    trainer.save_checkpoint(config.ckpt_path)
    trainer.test(model,data_mnist)


# https://github.com/neptune-ai/examples/blob/main/integrations-and-supported-tools/pytorch-lightning/scripts/Neptune_Pytorch_Lightning_more_options.py
# https://app.neptune.ai/common/pytorch-lightning-integration/e/PTL-15/source-code?path=source_code&file=complex.py&attribute=files&filePath=.

